use crate::{agent::Agent, game::Game, skill::Skill, strategy::Shape};
use rand::{prelude::*, seq::index::sample_weighted};
use std::{cell::Cell, cmp::Reverse};

pub struct World<'a, S: Skill> {
    /// 保持的大小
    size: usize,
    shape: &'a Shape<S>,
    /// 非空，包括分数（一般没用），最长是三倍大小
    agents: Vec<(Agent<S>, Cell<i32>)>,
}

impl<'a, S: Skill> World<'a, S> {
    pub fn new(size: usize, shape: &'a Shape<S>) -> Self {
        let rng = &mut thread_rng();
        let mut agents = Vec::with_capacity(size * 3);
        let stream = std::iter::repeat_with(|| (Agent::new_rand(shape, rng), Cell::new(0)));
        agents.extend(stream.take(size));
        Self {
            size,
            shape,
            agents,
        }
    }

    pub fn from_agents(agents: impl IntoIterator<Item = Agent<S>>, shape: &'a Shape<S>) -> Self {
        let mut agents: Vec<_> = agents.into_iter().map(|a| (a, Cell::new(0))).collect();
        let size = agents.len();
        agents.reserve_exact(size * 2);
        Self {
            size,
            shape,
            agents,
        }
    }

    pub fn agents(&self) -> impl Iterator<Item = &Agent<S>> {
        self.agents.iter().map(|(a, _)| a)
    }

    fn reproduce(&mut self) {
        let rng = &mut thread_rng();

        self.agents.shuffle(rng);
        self.agents.extend_from_within(0..self.size);
        self.agents.extend_from_within(0..self.size);
        let (orig, new) = self.agents.split_at_mut(self.size);
        let (mutate, blend) = new.split_at_mut(self.size);

        for agent in mutate {
            agent.0.mutate(self.shape, rng);
        }
        let skip = rng.gen_range(1..self.size);
        for (agent, other) in blend.iter_mut().zip(orig.iter().cycle().skip(skip)) {
            agent.0.blend(self.shape, &other.0, rng);
        }
    }

    fn race(&mut self, times: usize, max_round: usize) {
        let (mut rng, mut rng2) = (thread_rng(), thread_rng());

        for (_, score) in &self.agents {
            score.set(0); // 重置分数
        }
        self.agents.shuffle(&mut rng);

        for _ in 0..times {
            let n = rng.gen_range(1..self.agents.len());
            let games = self.agents.iter().zip(self.agents.iter().cycle().skip(n));
            for ((agent1, score1), (agent2, score2)) in games {
                let mut game = Game::new(
                    (agent1, self.shape, &mut rng),
                    (agent2, self.shape, &mut rng2),
                    false,
                );
                // 如果不分胜负，就不管他
                let result = (0..max_round).find_map(|_| game.round());
                if let Some(left) = result {
                    let (win, lose) = if left {
                        (score1, score2)
                    } else {
                        (score2, score1)
                    };
                    win.set(win.get() + 1);
                    lose.set(lose.get() - 1);
                }
            }
        }
        self.agents.sort_unstable_by_key(|(_, s)| Reverse(s.get()));
    }

    fn select(&mut self) {
        let dead = sample_weighted(
            &mut thread_rng(),
            self.agents.len(),
            |i| i as f64,
            self.agents.len() - self.size,
        )
        .expect("选择错误");
        for index in dead {
            self.agents[index].1.set(i32::MAX); // 假定没有这种分数
        }
        self.agents.retain(|s| s.1.get() != i32::MAX);
    }

    pub fn cycle(&mut self) {
        self.reproduce();
        self.race(10, 200);
        self.select();
    }

    pub fn best(&mut self) -> &Agent<S> {
        self.race(200, 500);
        &self.agents.first().expect("没有").0
    }
}
